"""
    StateVectors

Minimal StateVectors module for testing.
"""
module StateVectors

using LinearAlgebra
using SparseArrays
using StaticArrays
using Distributed
using FFTW
using MPI
using Printf

using ..GSICoreAnalysis: AbstractAnalysisConfig, AbstractStateVector, AnalysisVariables, default_analysis_variables, GridDimensions

export StateVector, allocate_state!

"""
    StateVector{T<:AbstractFloat}

Minimal state vector for testing with essential atmospheric variables.
"""
mutable struct StateVector{T<:AbstractFloat} <: AbstractStateVector{T}
    u::Array{T,3}         # Zonal wind [m/s]
    v::Array{T,3}         # Meridional wind [m/s] 
    t::Array{T,3}         # Temperature [K]
    q::Array{T,3}         # Specific humidity [kg/kg]
    ps::Matrix{T}         # Surface pressure [Pa]
    config::AbstractAnalysisConfig
    is_allocated::Bool
    
    function StateVector{T}(config::AbstractAnalysisConfig) where T
        sv = new{T}(
            Array{T,3}(undef, 0, 0, 0),
            Array{T,3}(undef, 0, 0, 0),
            Array{T,3}(undef, 0, 0, 0),
            Array{T,3}(undef, 0, 0, 0),
            Matrix{T}(undef, 0, 0),
            config,
            false
        )
        # Allocate memory immediately
        allocate_state!(sv)
        return sv
    end
end

StateVector(config::AbstractAnalysisConfig) = StateVector{Float64}(config)

"""
    allocate_state!(state::StateVector{T}) where T

Allocate memory for state vector variables based on grid dimensions.
"""
function allocate_state!(state::StateVector{T}) where T
    if state.is_allocated
        return state
    end
    
    # Get grid dimensions from config
    grid_size = state.config.grid_size
    nx, ny, nz = grid_size[1], grid_size[2], grid_size[3]
    
    # Allocate 3D fields with proper dimensions
    state.u = zeros(T, nx, ny, nz)
    state.v = zeros(T, nx, ny, nz)
    state.t = zeros(T, nx, ny, nz)
    state.q = zeros(T, nx, ny, nz)
    
    # Allocate 2D field
    state.ps = zeros(T, nx, ny)
    
    state.is_allocated = true
    
    println("StateVector allocated with dimensions: $(nx) x $(ny) x $(nz)")
    println("  Total 3D fields: $(4) with $(nx*ny*nz) points each")
    println("  Total 2D fields: $(1) with $(nx*ny) points each")
    
    return state
end

end # module StateVectors